#!/usr/env/bin python
# -*- coding: utf-8 -*-

# @Time    : 2018/8/23 19:17
# @Author  : yangdy
# @Email   : yangdy@egu360.com
# @File    : 后缀树.py
# @Software: PyCharm


class SuffixTree:
    class Node:
        def __init__(self, x, e):
            self.i = x
            self.e = e
            self.children = None
            self.brother = None
            self.end = True

    def __init__(self, s):
        self.s = s
        self.len = len(s)
        self.root = self.Node(0, self.len)

        def newNode(x, e=self.len):
            return self.Node(x, e)

        def add(nd=self.root, x=0):
            if not nd:
                self.root = newNode(x)
                return
            p = nd
            while p:
                if self.s[p.i] == self.s[x]:
                    break
                p = p.brother
            if not p:
                q = newNode(x)
                q.brother = nd.brother
                nd.brother = q
            else:
                ii = 0
                ll = p.e-p.i if p.e else self.len-p.i
                ll = ll if ll <= self.len-x else self.len-x
                while ii < ll:
                    if self.s[p.i+ii] != self.s[x+ii]:
                        break
                    ii += 1

                if ii >= ll and p.e-p.i < self.len-x:
                    add(p, x+ll)
                elif ii >= ll and p.e-p.i == self.len-x:
                    p.end = True
                else:
                    q = newNode(p.i+ii)
                    q.e = p.e
                    p.e = p.i+ii
                    q.children = p.children
                    p.children = q
                    if ii < ll:
                        r = newNode(x+ii)
                        r.brother = q
                        p.children = r
                        p.end = False

        add()
        self.display()
        for i in range(1, self.len):
            add(x=i)

    def display(self):
        def display(nd, dp=0, ss=''):
            if not nd:
                return
            if nd.brother:
                tag = '├─'
            else:
                tag = '└─'
            end = '(end)' if nd.end else ''
            pre = ss + '|' if ss else ''
            print('%s%s%d:%s' % ('  '*dp, tag, nd.i, pre+self.s[nd.i:nd.e]+end))
            p = nd.children
            while p:
                display(p, dp+1, ss+self.s[nd.i:nd.e])
                p = p.brother
        p = self.root
        while p:
            display(p)
            p = p.brother


def main():
    st = SuffixTree('acbbacb')
    print('ok')
    st.display()
    print('------end-------')
    print('start...')
    st1 = SuffixTree('abceabd')
    st1.display()
    print('end..')


if __name__ == '__main__':
    main()
